import json
import sys
from typing import List, Optional, Dict, Callable,Any, Union  # 添加 List 导入
from functools import wraps
import time
from termcolor import colored
import threading
from sqlmodel import SQLModel, Field, Relationship, Column
from typing import List, Dict, Optional, Union
from sqlalchemy import ARRAY, Float
from sqlalchemy.dialects.postgresql import JSONB
from pydantic import PrivateAttr
from typing import ClassVar
from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.orm import Session
from sqlalchemy import func
from sqlmodel import create_engine, Session
import traceback
import warnings
from sqlalchemy import text
import os
import subprocess#type:ignore
warnings.filterwarnings("ignore", category=DeprecationWarning)
AFFINE_TEST_BASE_PATH = r"D:\CYBERLIFE_\NoteAgentTree3\testaffBE"
SOURCE_ROOT=r"D:\CYBERLIFE_\NoteAgentTree3\testIndexroot"
def is_docker() -> bool:
    """
    Check if the code is running inside a Docker container.
    Returns:
        bool: True if running in Docker, False otherwise
    """
    path = "/.dockerenv"
    if os.path.exists(path):
        return True
    try:
        with open("/proc/self/cgroup", "r") as f:
            for line in f:
                if "docker" in line:
                    return True
    except FileNotFoundError:
        pass
    return False
def docker_init():
    if is_docker():
        subprocess.run(['service', 'postgresql', 'start'])
        SQLModel.metadata.create_all(engine,checkfirst=True)
if is_docker():
    DBURL="postgresql://noteagent:noteagent@localhost:5432/noteagent?sslmode=disable"
else:
    DBURL="postgresql://postgres:142857@localhost:5432/postgres?sslmode=disable"
# 创建全局的 engine（它是线程安全的）
engine = create_engine(DBURL, echo=False)
def get_session():
    return Session(engine) 
def format_error(e: Exception) -> Dict[str, Any]:
    """ 格式化异常为标准错误字典 """
    tbstr=traceback.format_exc()+"\n"
    print(colored(tbstr,"red"))
    if hasattr(e, 'args') and e.args and isinstance(e.args[0], dict):
        return e.args[0]
    return {"success": -1, "message": f"(error type unknown){tbstr+str(e)}"}
def escapeLn(string: Union[str, dict]) -> str:
    if isinstance(string, dict):
        return json.dumps(string, ensure_ascii=False).replace('\\n', '\n')
    elif isinstance(string, str):   
        decoded_str = string
        # .encode('utf-8').decode('unicode_escape')
        return decoded_str.replace('\\n', '\n')
    else:
        print(f"Invalid input type: {type(string)}")
        return str(string)
def rollback_on_exception(func):
    """
    统一的数据库异常处理装饰器：
    - 检查含有'int'的参数是否为有效整数
    - 若被装饰的方法返回 {"success": -1} 或更小的值，则自动 rollback
    - 若抛出异常，也自动 rollback 并返回 self.err.INTERNAL_EXCEPTION
    """
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        # 获取函数参数名
        arg_names = func.__code__.co_varnames[:func.__code__.co_argcount]
        
        # 检查位置参数
        for i, (name, value) in enumerate(zip(arg_names[1:], args)):  # 跳过self
            if 'int' in name.lower():
                try:
                    int(str(value))
                except ValueError:
                    return self.err.INVALID_PARAM(f"{func.__name__}.{name} receives an int but {value} is not an int")
        
        # 检查关键字参数
        for name, value in kwargs.items():
            if 'int' in name.lower():
                try:
                    int(str(value))
                except ValueError:
                    return self.err.INVALID_PARAM(f"{func.__name__}.{name} receives an int but {value} is not an int")
                    
        try:
            rtn_dict = func(self, *args, **kwargs)
            if isinstance(rtn_dict, dict) and int(rtn_dict.get("success", -1)) < 0:
                self.session.rollback()
            return rtn_dict
        except Exception as e:
            self.session.rollback()
            return self.err.INTERNAL_EXCEPTION(e)
    return wrapper
class ErrorCode:
    session:Session
    def __init__(self,session:Session):
        ErrorCode.session=session
    @staticmethod
    def NO_PERMISSION(message:str="no permission") -> Dict[str, Union[str,int]]:
        return {"success": -1, "message": message}  # 用户没有权限
    @staticmethod
    def MAPFRAME_NOT_FOUND(mapType: str,intId:int=0) -> Dict[str, Union[str,int]]:
        list_affmap=ErrorCode.session.query(affmap).all()
        for item in list_affmap:
            print(item)
        addstr=""
        if str(intId)=="1":
            addstr+="You may be trying to manipulate the root directory, upon which only creating spaces is supported"
        return {"success": -2, "message": f"{mapType} not found. {addstr}"}      # 资源未找到
        
    @staticmethod
    def INTERNAL_EXCEPTION(e:Exception) -> Dict[str, Union[str,int]]:
        tbstr=traceback.format_exc()  # 打印完整的异常堆栈信息
        print(colored(tbstr,"red"))
        return {"success": -3, "message": f"Internal exception: \n{tbstr}\n{e}"}      # 内部异常
    @staticmethod
    def NODE_NOT_FOUND(parentType:str,intId:int) -> Dict[str, Union[str,int]]:
        return {"success": -4, "message": f"{parentType} intId={intId} not found"}          # 父节点未找到
    @staticmethod
    def TYPE_NOT_MATCH(type:str) -> Dict[str, Union[str,int]]:
        return {"success": -5, "message": f"type not match: {type}"}          # 类型不匹配
    @staticmethod
    def INVALID_PARAM(param:str) -> Dict[str, Union[str,int]]:
        return {"success": -6, "message": f"invalid param: {param}"}          # 参数无效
    @staticmethod
    def UNKNOWN_TOOL_CALL(name:str) -> Dict[str, Union[str,int]]:
        return {"success": -7, "message": f"unknown tool call: {name}"}          # 未知工具调用
    @staticmethod
    def USER_REJECTED(name:str) -> Dict[str, Union[str,int]]:
        return {"success": -8, "message": f"user rejected: {name}"}          # 用户拒绝
    @staticmethod
    def SUCCESS(id: int) -> Dict[str, Union[str,int]]:
        if id<0:
            return {"success": id, "message": "failed"}  # 失败
        else:
            return {"success": id, "message": "success"}  # 成功
    @staticmethod
    def IF_SUCCESS_COMMIT(i_dictOrId: Union[Dict[str, Union[str,int]],int],allowCommit:bool=True) -> Dict[str, Union[str,int]]:
        """
        获取i_dict的success，>=0则commit()，然后返回i_dict.供需要commit的crudUser函数使用
        """
        sys.stdout.flush()
        if isinstance(i_dictOrId,int):
            rtndict=ErrorCode.SUCCESS(i_dictOrId)
        else:
            rtndict=i_dictOrId
        i_id=int(rtndict.get("success",-1))
        if i_id<0:
            return rtndict
        if allowCommit:
            ErrorCode.session.commit()
        return rtndict
    @staticmethod
    def IF_SUCCESS_COMMIT_A2I(a_dict: Dict[str, Union[str,int]],i_dict: Dict[str, Union[str,int]]) -> Dict[str, Union[str,int]]:
        """
        获取i_dict的success，并赋值给a_dict的success,然后返回a_dict.供需要commit的crudAgent函数使用
        """
        sys.stdout.flush()
        a_id=int(a_dict.get("success",-1))
        if a_id<0:
            return a_dict
        i_id=int(i_dict.get("success",-1))
        assert i_id>=0
        a_dict["success"]=i_id
        # print("before commit",ErrorCode.session.query(affmap).all())
        ErrorCode.session.commit()
        ErrorCode.session.flush()
        # print("affmap after commit",ErrorCode.session.query(affmap).all())
        return a_dict
def trace_info(color='blue',Print=False):
    """
    装饰器：打印函数名称、参数名称及传入值，以及函数返回值
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 获取函数参数名
            arg_names = func.__code__.co_varnames[:func.__code__.co_argcount]
            
            # 安全配对参数名和传入值
            paired_args = []
            for i, arg_value in enumerate(args):
                if i < len(arg_names):
                    name = arg_names[i]
                    if name not in ["self", "cls"]:
                        if name.endswith("permission"):
                            paired_args.append(f"{name}={int(arg_value):08b}")
                        else:
                            paired_args.append(f"{name}={str(arg_value)[0:100]}")
                else:
                    # 处理超出定义参数的额外位置参数
                    paired_args.append(f"*extra_arg_{i-len(arg_names)+1}={str(arg_value)[0:100]}")

            # 拼接输出字符串：函数名 + 参数列表
            info_str = func.__name__ + "(" + ", ".join(paired_args) + ")"
            if Print:
                print(colored(info_str, color))
            
            result = func(*args, **kwargs)

            if isinstance(result, dict):
                # 先将 dict 转成字符串，再把其中的 \n 替换为真正的换行
                raw_string = json.dumps(result, ensure_ascii=False)
                raw_string = raw_string.replace('\\n', '\n')
                if Print:
                    print(colored(f"{func.__name__}:return:{raw_string[0:100]}", color))
            elif Print:
                print(colored(f"{func.__name__}:return:{result[0:100] if isinstance(result,str) else result}", color))
            return result

        return wrapper
    return decorator
def timeit(color='magenta'):
    """装饰器：测量函数执行时间，并允许指定输出颜色"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            print(colored(f"{func.__name__} time:{(end_time - start_time):.3f}s", color))
            sys.stdout.flush()
            return result
        return wrapper
    return decorator
class BackCount:
    """
    倒计时对象，支持设置倒计时和回调函数（单例模式）
    """
    _instance = None
    _timer: Optional[threading.Timer] = None  # 添加类型注解
    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(BackCount, cls).__new__(cls)
            cls._instance._timer = None
            cls._instance._callback = None
        return cls._instance
    def start_countdown(self, seconds: int, callback: Callable[[], Any]):
        """
        开始倒计时
        :param seconds: 倒计时秒数
        :param callback: 倒计时结束时执行的回调函数
        如果倒计时正在运行，则不执行
        """
        if self._timer:
            self.clear_countdown()
        self._callback = callback
        self._timer = threading.Timer(seconds, self._execute_callback)
        self._timer.start()
    def clear_countdown(self):
        """
        如果倒计时正在运行，则清除当前倒计时，并阻止回调函数的执行
        """
        if self._timer:
            self._timer.cancel()
            self._timer = None
            self._callback = None
    def _execute_callback(self):
        """
        执行回调函数
        """
        if self._callback:
            self._callback()
        self._timer = None
        self._callback = None
priority_map = {
    'space': 1,
    'doc': 2,
    'h1': 3, 'h2': 4, 'h3': 5, 'h4': 6, 'h5': 7, 'h6': 8,  # 扩展所有标题级别
    'ul': 9, 'ol': 9, 'li': 10, 'p': 11,
    'blockquote': 12,  # 引用块
    'pre': 13, 'code': 14,  # 代码块和内联代码
    'hr': 15,  # 水平线
    'img': 16,  # 图片
    'a': 17,   # 链接
    'br': 18,  # 换行
    'table': 19,
    'tr': 20,
    'td': 21,
    'th': 22,
    'thead': 23,
    'tbody': 24,
    'span': 25,
    'div': 26,
    'section': 27,
    'strong': 28,
    'em': 29,
    'all': 101
}

def get_previous_priority(item_key):
    keys = list(priority_map.keys())  # 直接获取键的列表
    if item_key in keys:
        current_index = keys.index(item_key)
        previous_key = keys[current_index - 1] if current_index > 0 else None
        return previous_key
    else:
        return "space"
class autoapprove_oper(SQLModel, table=True):
    __tablename__ = "autoapprove_oper"  # 添加明确的表名声明
    user: str = Field(primary_key=True)
    operation: str = Field(primary_key=True)
    def __init__(self, user:str, operation:str):
        self.user = user
        self.operation = operation

class cache_emb(SQLModel, table=True):
    __tablename__ = "cache_emb"  # 添加明确的表名声明
    id:int = Field(default=None, primary_key=True, sa_column_kwargs={"autoincrement": True})
    content:str = Field(default="")
    vec:List[float] = Field(default_factory=list, sa_type=ARRAY(Float))  # type: ignore
    def __init__(self, content:str, vec:List[float]):
        self.content = content
        self.vec = vec
class affmap(SQLModel, table=True):
    __tablename__ = "affmap"  # 添加明确的表名声明
    a_docId:str = Field(default="",index=True)
    a_spaceId:str = Field(default="",index=True)
    intId:int = Field(default=-1,index=True,primary_key=True)
    def __init__(self, a_docId:str="", a_spaceId:str="", intId:int=-1):
        self.a_docId = a_docId
        self.a_spaceId = a_spaceId
        self.intId = intId
class fnode(SQLModel, table=True):  # 添加 table=True 以启用表映射
    __tablename__ = "fnode"  # 添加明确的表名声明
    """
    fnode 结构体（以 Pydantic 模型形式表达）:
    - id: 节点ID
    - content: 节点内容
    - node_type: 节点类型 ("space","doc","p","h1","h2" 等)
    - summary: 由 LLM 或者其他手段生成的概述。
    - sum_vec: summary 的向量表示
    - content_vec: content 的向量表示
    - parent: 上一个节点
    - children: 子节点列表
    - permission: {userId: int 权限掩码}
    - start_line: 起始行号
    - end_line: 结束行号
    - _extra_info: 额外信息
    - _sim_score: 相似度分数
    - level: 节点层级
    """
    _fnode_counter: ClassVar[int] = 0
    id: int = Field(default=None, primary_key=True, sa_column_kwargs={"autoincrement": True})
    content: str = ""
    node_type: str = Field(default="space")
    summary: str = ""
    sum_vec: List[float] = Field(default_factory=list, sa_type=ARRAY(Float))  # type: ignore
    content_vec: List[float] = Field(default_factory=list, sa_type=ARRAY(Float))  # type: ignore
    parent_id: Optional[int] = Field(
        default=None,
        foreign_key="fnode.id"  # 指向自身的主键
    )
    children: List["fnode"] = Relationship(
        # back_populates="parent"
        )
    parent: Optional["fnode"] = Relationship(
        # back_populates="children",
        sa_relationship_kwargs={"remote_side": "fnode.id","overlaps":"children"}
    )
    permission: Dict[str, int] = Field(
        default_factory=MutableDict,
        sa_column=Column(MutableDict.as_mutable(JSONB))
    )
    start_line: int = Field(default=0)
    end_line: int = Field(default=0)
    _sim_score: List[float] = PrivateAttr(default_factory=list)
    _extra_info: str = PrivateAttr(default="")
    level: int = Field(default=0)
    mdSeq: int = Field(default=0)
    def __init__(
        self,
        content: str,
        node_type: str,
        summary: str = "",
        sum_vec: Optional[List[float]] = None,
        content_vec: Optional[List[float]] = None,
        parent_node: Optional['fnode'] = None,
        next_nodes: Optional[List['fnode']] = None,
        permission: Optional[Dict[str, int]] = None,
        start_line: int = 0,
        end_line: int = 0,
        _extra_info: str = "",
        _sim_score: Optional[List[float]] = None,
        level: int = 0,
        mdSeq: int = -1,
    ):
        # 使用普通类变量进行计数自增
        # max_id = ftree.session.query(func.max(fnode.id)).scalar() or 0
        # type(self)._fnode_counter = max_id
        if permission is None:
            permission={}
        super().__init__(
            # id=type(self)._fnode_counter,  # 用计数器赋给 id
            id=None,
            content=content,
            node_type=node_type,
            summary=summary,
            sum_vec=sum_vec if sum_vec else [],
            content_vec=content_vec if content_vec else [],
            parent_id=parent_node.id if parent_node else None,  # 引用父节点的 id
            children=next_nodes if next_nodes else [],
            permission=permission, # 确保 permission 始终是字典
            start_line=start_line,
            end_line=end_line,
            level=level,
            mdSeq=mdSeq,
        )
        # 将原本 self._sim_score = _sim_score or [] 改为 setattr
        setattr(self, '_sim_score', _sim_score or [])

        # 将原本 self._extra_info = _extra_info 改为 setattr
        setattr(self, '_extra_info', _extra_info)

        self.parent = parent_node
        self.permission = permission
    def __hash__(self):
        return hash(self.id)
    def __eq__(self, other):
        return (
            isinstance(other, fnode)
            and self.id == other.id
        )

def reinitDB():
    """
    删除并重新初始化数据库中的所有表
    """
    # with engine.connect() as conn:
    #     # 删除表
    #     SQLModel.metadata.drop_all(engine, tables=[affmap.__table__, fnode.__table__, autoapprove_oper.__table__])
        
    #     # 重置自增序列（PostgreSQL语法）
    #     conn.execute(text("ALTER SEQUENCE IF EXISTS fnode_id_seq RESTART WITH 1"))
        
    #     # 创建表
    #     SQLModel.metadata.create_all(engine)
    # engine.dispose()
    SQLModel.metadata.drop_all(engine)
    SQLModel.metadata.create_all(engine,checkfirst=True)
    
# 使用示例
if __name__ == "__main__":
    if input("reinitDB? this operation will clear all tables (y/n)")=="y":
        reinitDB()


